package com.mical.demo.controller;


import com.mical.demo.respository.MessageRepository;
import com.mical.demo.security.JwtTokenUtil;
import com.mical.demo.service.AuthService;
import com.mical.demo.service.DifyChatDefaultImpl;
import io.github.guoshiqiufeng.dify.chat.dto.request.ChatMessageSendRequest;
import io.github.guoshiqiufeng.dify.chat.dto.response.ChatMessageSendResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;

import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.http.MediaType;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;

import org.springframework.web.bind.annotation.*;

import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
@RestController
@RequestMapping("/api/chat")
public class ChatController {
    private static final Logger logger = LoggerFactory.getLogger(ChatController.class);
    private static final long SSE_TIMEOUT = TimeUnit.MINUTES.toMillis(5);
    private final MessageRepository messageRepository;
    private final AuthService authService;
    private final DifyChatDefaultImpl chatService;
    private final JwtTokenUtil jwtTokenUtil;

    @Autowired
    @Qualifier("sseTaskExecutor")
    private AsyncTaskExecutor asyncTaskExecutor;

    public ChatController(MessageRepository messageRepository, AuthService authService,
                          DifyChatDefaultImpl chatService,
                          JwtTokenUtil jwtTokenUtil) {
        this.messageRepository = messageRepository;
        this.authService = authService;
        this.chatService = chatService;
        this.jwtTokenUtil = jwtTokenUtil;
    }

    @PreAuthorize("hasAuthority('USER')")
    @PostMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter handleChatStream(@RequestBody ChatMessageSendRequest request,
                                       Authentication auth) {
        // 1. 初始化上下文
        SecurityContext context = SecurityContextHolder.createEmptyContext();
        context.setAuthentication(auth);
        SecurityContextHolder.setContext(context);

        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT);
        String userId = request.getUserId();

        // 保存用户输入消息到数据库
        messageRepository.saveMessage(request.getContent(),Integer.parseInt(userId),1);

        // 2. 使用连接状态跟踪器
        ConnectionTracker tracker = new ConnectionTracker();

        // 3. 异步处理流
        asyncTaskExecutor.execute(() -> {
            try {
                Disposable subscription = chatService.sendChatMessageStream(request)
                        .publishOn(Schedulers.boundedElastic())
                        .doOnNext(response -> {
                            if (response.getAnswer() != null) {
                                tracker.appendResponse(response.getAnswer());
                            }
                        })
                        .takeWhile(item -> tracker.isActive())
                        .doFinally(signal -> SecurityContextHolder.clearContext())
                        .subscribe(
                                response -> sendSSEEvent(emitter, response, tracker),
                                error -> handleError(emitter, error, userId, tracker),
                                () -> completeEmitter(emitter, request, userId, tracker)
                        );

                tracker.setSubscription(subscription);
                setupEmitterCallbacks(emitter, request, userId, tracker);

            } catch (Exception e) {
                logger.error("初始化流失败 - 用户: {}: {}", userId, e.getMessage());
                completeEmitterWithError(emitter, e);
                SecurityContextHolder.clearContext();
            }
        });

        return emitter;
    }

    // 连接状态跟踪器（内部类）
    private static class ConnectionTracker {
        private final AtomicBoolean active = new AtomicBoolean(true);
        private final AtomicBoolean completed = new AtomicBoolean(false);
        private final StringBuffer responseBuffer = new StringBuffer();
        private Disposable subscription;

        public boolean isActive() {
            return active.get() && !completed.get();
        }

        public void markCompleted() {
            completed.set(true);
            active.set(false);
        }

        public void appendResponse(String content) {
            responseBuffer.append(content);
        }
        public String getFullResponse() {
            return responseBuffer.toString();
        }
        public void setSubscription(Disposable subscription) {
            this.subscription = subscription;
        }

        public void dispose() {
            if (subscription != null && !subscription.isDisposed()) {
                subscription.dispose();
            }
        }
    }

    private void sendSSEEvent(SseEmitter emitter, ChatMessageSendResponse response,
                              ConnectionTracker tracker) {
        if (!tracker.isActive()) return;

        try {
            emitter.send(SseEmitter.event()
                    .id(response.getConversationId())
                    .name("message")
                    .data(response));
        } catch (IOException e) {
            logger.debug("客户端断开连接 - 会话ID: {}", response.getConversationId());
            tracker.markCompleted();
            tryCompleteEmitter(emitter);
        } catch (IllegalStateException e) {
            logger.debug("Emitter已完成状态异常");
            tracker.markCompleted();
        }
    }

    private void handleError(SseEmitter emitter, Throwable error, String userId,
                             ConnectionTracker tracker) {
        logger.error("流处理错误 - 用户: {}: {}", userId, error.getMessage());
        tracker.markCompleted();
        completeEmitterWithError(emitter, error);
    }

    private void completeEmitter(SseEmitter emitter, ChatMessageSendRequest request,
                                 String userId, ConnectionTracker tracker) {
        try {
            // 保存完整的AI输出消息
            String fullResponse = tracker.getFullResponse();
            if (fullResponse != null && !fullResponse.isEmpty()) {
                messageRepository.saveMessage(fullResponse,Integer.parseInt(userId),2);
            }
        } finally {
            tracker.markCompleted();
            stopDifyStream(request, userId);
            tryCompleteEmitter(emitter);
            logger.debug("流正常完成 - 用户: {}", userId);
        }
    }

    private void setupEmitterCallbacks(SseEmitter emitter, ChatMessageSendRequest request,
                                       String userId, ConnectionTracker tracker) {
        emitter.onTimeout(() -> {
            logger.debug("SSE超时 - 用户: {}", userId);
            cleanup(emitter, request, userId, tracker);
        });

        emitter.onCompletion(() -> {
            logger.debug("客户端主动断开 - 用户: {}", userId);
            cleanup(emitter, request, userId, tracker);
        });

        emitter.onError(error -> {
            logger.error("Emitter错误 - 用户: {}: {}", userId, error.getMessage());
            cleanup(emitter, request, userId, tracker);
        });
    }

    private void cleanup(SseEmitter emitter, ChatMessageSendRequest request,
                         String userId, ConnectionTracker tracker) {
        tracker.markCompleted();
        stopDifyStream(request, userId);
        tracker.dispose();
        tryCompleteEmitter(emitter);
    }

    private void tryCompleteEmitter(SseEmitter emitter) {
        try {
            emitter.complete();
        } catch (Exception e) {
            logger.trace("Emitter完成异常: {}", e.getMessage());
        }
    }

    private void completeEmitterWithError(SseEmitter emitter, Throwable error) {
        try {
            emitter.completeWithError(error);
        } catch (Exception e) {
            logger.warn("Emitter错误完成异常: {}", e.getMessage());
        }
    }

    private void stopDifyStream(ChatMessageSendRequest request, String userId) {
        try {
            if (request.getConversationId() != null && request.getApiKey() != null) {
                chatService.stopMessagesStream(
                        request.getApiKey(),
                        request.getConversationId(),
                        userId
                );
            }
        } catch (Exception e) {
            logger.error("停止Dify流失败 - 用户: {}: {}", userId, e.getMessage());
        }
    }
}


//@RestController
//@RequestMapping("/api/chat")
//public class ChatController {
//    private static final Logger logger = LoggerFactory.getLogger(ChatController.class);
//    private static final long SSE_TIMEOUT = TimeUnit.MINUTES.toMillis(5); // 5 minutes
//    //   private final MessageRepository messageRepository;
//    private final AuthService authService;
//    @Autowired
//    @Qualifier("applicationTaskExecutor")
//    private AsyncTaskExecutor asyncTaskExecutor;
//    private final DifyChatDefaultImpl chatService;
//    private final JwtTokenUtil jwtTokenUtil;
//    private final AtomicBoolean connectionActive = new AtomicBoolean(true);
//
//    public ChatController(
//            //  MessageRepository messageRepository,
//            AuthService authService,
//            DifyChatDefaultImpl chatService,
//            JwtTokenUtil jwtTokenUtil) {
//        // this.messageRepository = messageRepository;
//        this.authService = authService;
//
//        this.chatService = chatService;
//        this.jwtTokenUtil = jwtTokenUtil;
//    }
//
//    @PreAuthorize("hasAuthority('USER')")
//    @PostMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
//    public SseEmitter handleChatStream(@RequestBody ChatMessageSendRequest request, Authentication auth) {
//        // 首先保存用户提问
////        if (request.getContent() != null && !request.getContent().isEmpty()) {
////            messageRepository.saveMessage(request.getContent(),
////                    Integer.parseInt(request.getUserId()), 1); // 1代表提问
////        }
//        SecurityContext context = SecurityContextHolder.createEmptyContext();
//        context.setAuthentication(auth);
//        SecurityContextHolder.setContext(context);
//
//        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT);
//        connectionActive.set(true);
//
//        String userId = request.getUserId();
//        Disposable[] subscriptionHolder = new Disposable[1];
//        //字符拼接
//        StringBuilder accumulatedResponse = new StringBuilder();
//        asyncTaskExecutor.execute(() -> {
//            try {
//                Flux<ChatMessageSendResponse> responseFlux = chatService.sendChatMessageStream(request)
//                        .doOnNext(response -> {
//                            // 累积响应内容
//                            if (response.getAnswer() != null) {
//                                accumulatedResponse.append(response.getAnswer());
//                            }
//                            //                           else
////                            {
////                                // 仅在流结束时保存完整的累积响应
////                                if (!accumulatedResponse.isEmpty()) {
////                                    messageRepository.saveMessage(
////                                            accumulatedResponse.toString(),
////                                            Integer.parseInt(userId),
////                                            2 // 2代表回复
////                                    );
////                                }
////                            }
//                        })
//                        .doOnCancel(() -> {
//                            logger.debug("Flux was cancelled for user {}", userId);
////                            if (!accumulatedResponse.isEmpty()) {
////                                messageRepository.saveMessage(
////                                        accumulatedResponse.toString(),
////                                        Integer.parseInt(userId),
////                                        2 // 2代表回复
////                                );
////                            }
//                            cleanupConnection(emitter, request, userId, subscriptionHolder[0]);
//                        })
//                        .doOnTerminate(() -> {
//                            logger.debug("Flux terminated for user {}", userId);
//
//                            cleanupConnection(emitter, request, userId, subscriptionHolder[0]);
//                        })
//                        .doOnError(error -> {
//                            logger.error("Error in SSE stream for user {}: {}", userId, error.getMessage());
//                            cleanupConnection(emitter, request, userId, subscriptionHolder[0]);
//                        })
//                        .takeWhile(item -> connectionActive.get());
//                subscriptionHolder[0] = responseFlux.subscribe(
//                        response -> {
//                            if (connectionActive.get()) {
//                                sendSSEEvent(emitter, response);
//
//                            }
//                        },
//                        error -> {
//                            logger.error("Subscription error for user {}: {}", userId, error.getMessage());
//                            // 错误时尝试保存已累积的内容
//                            //                         if (!accumulatedResponse.isEmpty()) {
////                                messageRepository.saveMessage(
////                                        accumulatedResponse.toString(),
////                                        Integer.parseInt(userId),
////                                        2 // 2代表回复
////                                );
//                            //                          }
//                            completeEmitterWithError(emitter, error);
//                        },
//                        () -> {
//                            logger.debug("Subscription completed for user {}", userId);
//                            completeEmitter(emitter, request, userId);
//
//                        }
//                );
//
//                setupEmitterCallbacks(emitter, request, userId, subscriptionHolder[0]);
//
//            } catch (Exception e) {
//                logger.error("Error initializing chat stream for user {}: {}", userId, e.getMessage());
//                completeEmitterWithError(emitter, e);
//            } finally {
//                SecurityContextHolder.clearContext();
//            }
//        });
//
//        return emitter;
//    }
//
//    private boolean isEmitterActive(SseEmitter emitter) {
//        try {
//            emitter.send(SseEmitter.event().comment("ping"));
//            return true;
//        } catch (IOException | IllegalStateException e) {
//            return false;
//        }
//    }
//
//    private void sendSSEEvent(SseEmitter emitter, ChatMessageSendResponse response) {
//        if (!isEmitterActive(emitter) || !connectionActive.get()) {
//            return;
//        }
//
//        try {
//            emitter.send(SseEmitter.event()
//                    .id(response.getConversationId())
//                    .name("message")
//                    .data(response));
//        } catch (IOException e) {
//            logger.debug("Client disconnected, stopping SSE stream");
//            connectionActive.set(false);
//            try {
//                emitter.complete();
//            } catch (Exception ex) {
//                logger.debug("Error completing emitter: {}", ex.getMessage());
//            }
//        }
//    }
//
//    private void completeEmitter(SseEmitter emitter, ChatMessageSendRequest request, String userId) {
//        try {
//            stopDifyStream(request, userId);
//            emitter.complete();
//            logger.debug("SSE stream completed normally for user {}", userId);
//        } catch (Exception e) {
//            logger.warn("Error completing SSE emitter for user {}: {}", userId, e.getMessage());
//        }
//    }
//
//    private void completeEmitterWithError(SseEmitter emitter, Throwable error) {
//        try {
//            emitter.completeWithError(error);
//            logger.error("SSE stream completed with error: {}", error.getMessage());
//        } catch (Exception e) {
//            logger.warn("Error completing SSE emitter with error: {}", e.getMessage());
//        }
//    }
//
//    private void cleanupConnection(SseEmitter emitter, ChatMessageSendRequest request,
//                                   String userId, Disposable subscription) {
//        connectionActive.set(false);
//        stopDifyStream(request, userId);
//        if (subscription != null && !subscription.isDisposed()) {
//            subscription.dispose();
//        }
//        try {
//            emitter.complete();
//        } catch (Exception e) {
//            logger.debug("Error completing emitter during cleanup: {}", e.getMessage());
//        }
//    }
//
//    private void setupEmitterCallbacks(SseEmitter emitter, ChatMessageSendRequest request,
//                                       String userId, Disposable subscription) {
//        emitter.onTimeout(() -> {
//            logger.debug("SSE timeout detected for user {}", userId);
//            cleanupConnection(emitter, request, userId, subscription);
//        });
//
//        emitter.onCompletion(() -> {
//            logger.debug("SSE client disconnected (completion) for user {}", userId);
//            cleanupConnection(emitter, request, userId, subscription);
//        });
//
//        emitter.onError(error -> {
//            logger.error("SSE error occurred for user {}: {}", userId, error.getMessage());
//            cleanupConnection(emitter, request, userId, subscription);
//        });
//    }
//
//    private void stopDifyStream(ChatMessageSendRequest request, String userId) {
//        if (request.getConversationId() != null && request.getApiKey() != null) {
//            try {
//                chatService.stopMessagesStream(
//                        request.getApiKey(),
//                        request.getConversationId(),
//                        userId
//                );
//                logger.debug("Successfully stopped Dify stream for user {}", userId);
//            } catch (Exception e) {
//                logger.error("Error stopping Dify stream for user {}: {}", userId, e.getMessage());
//            }
//        }
//    }
//}